import torch

class SimBA:

    def __init__(self, neter):

        self.neter = neter

    def SimBA_attack(
        self,
        x,
        y,
        max_queries=5000,
        epsilon=0.05,
        targeted=False, 
    ):

        n_dims = x.view(1, -1).size(1)
        perm = torch.randperm(n_dims)

        last_prob = self.get_probs(x, y)
        query_count = 0

        for i in range(n_dims):
            
            if query_count > max_queries:
                break
            
            diff = torch.zeros(n_dims)
            diff[perm[i]] = epsilon
            left_prob = self.get_probs((x - diff.view(x.size())).clamp(0, 1), y)
            query_count += 1
            if targeted != (left_prob < last_prob):
                x = (x - diff.view(x.size())).clamp(0, 1)
                last_prob = left_prob
            else:
                right_prob = self.get_probs((x + diff.view(x.size())).clamp(0, 1), y)
                query_count += 1
                if targeted != (right_prob < last_prob):
                    x = (x + diff.view(x.size())).clamp(0, 1)
                    last_prob = right_prob
            
            if targeted:
                # target attack
                flag, index = self.neter.isTargetAttckSuccessful(x, torch.zeros_like(x), y)
                if flag:
                    print('Last probability: {:.2f}%'.format(last_prob * 100))
                    return flag, x
            else:
                # untarget attack
                flag, index = self.neter.isUntargetAttackSuccessful(x, torch.zeros_like(x), y)
                if flag:
                    print('Last probability: {:.2f}%'.format(last_prob * 100))
                    return flag, x

            if i % 100 == 0:
                print('Iter: {}, Prob: {}'.format(i, last_prob))

        return False, None
    
    def get_probs(self, x, y):

        return self.neter.get_prediction_score(x)[0, y]